import logging
import os
import time
import mindspore as ms
import cv2
import numpy as np
#import torch
from mindspore import nn
import yaml
from matplotlib import colors
from matplotlib import pyplot as plt
#from torch import Tensor, nn
#from torch.utils.data import ConcatDataset

class CharsetMapper(object):
    """A simple class to map ids into strings.

    It works only when the character set is 1:1 mapping between individual
    characters and individual ids.
    """

    def __init__(self,
                 filename='',
                 max_length=30,
                 null_char=u'\u2591'):
        """Creates a lookup table.

        Args:
          filename: Path to charset file which maps characters to ids.
          max_sequence_length: The max length of ids and string.
          null_char: A unicode character used to replace '<null>' character.
            the default value is a light shade block '░'.
        """
        self.null_char = null_char
        self.max_length = max_length

        self.label_to_char = self._read_charset(filename)
        self.char_to_label = dict(map(reversed, self.label_to_char.items()))
        self.num_classes = len(self.label_to_char)
 
    def _read_charset(self, filename):
        """Reads a charset definition from a tab separated text file.

        Args:
          filename: a path to the charset file.

        Returns:
          a dictionary with keys equal to character codes and values - unicode
          characters.
        """
        import re
        pattern = re.compile(r'(\d+)\t(.+)')
        charset = {}
        self.null_label = 0
        charset[self.null_label] = self.null_char
        with open('/home/data4/zyh/ABINet/data/charset_36.txt', 'r') as f:
            for i, line in enumerate(f):
                m = pattern.match(line)
                assert m, f'Incorrect charset file. line #{i}: {line}'
                label = int(m.group(1)) + 1
                char = m.group(2)
                charset[label] = char
        return charset

    def trim(self, text):
        assert isinstance(text, str)
        # print(self.null_char)
        return text.replace(self.null_char, '')

    def get_text(self, labels, length=None, padding=True, trim=False):
        """ Returns a string corresponding to a sequence of character ids.
        """
        length = length if length else self.max_length
        labels = [int(l) if isinstance(l, ms.Tensor ) else int(l) for l in labels]
        if padding:
            labels = labels + [self.null_label] * (length-len(labels))
        text = ''.join([self.label_to_char[label] for label in labels])
        if trim: text = self.trim(text)
        return text

    def get_labels(self, text, length=None, padding=True, case_sensitive=False):
        """ Returns the labels of the corresponding text.
        """
        length = length if length else self.max_length
        if padding:
            text = text + self.null_char * (length - len(text))
        if not case_sensitive:
            text = text.lower()
        labels = [self.char_to_label[char] for char in text]
        return labels

    def pad_labels(self, labels, length=None):
        length = length if length else self.max_length

        return labels + [self.null_label] * (length - len(labels))

    @property
    def digits(self):
        return '0123456789'

    @property
    def digit_labels(self):
        return self.get_labels(self.digits, padding=False)

    @property
    def alphabets(self):
        all_chars = list(self.char_to_label.keys())
        valid_chars = []
        for c in all_chars:
            if c in 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ':
                valid_chars.append(c)
        return ''.join(valid_chars)

    @property
    def alphabet_labels(self):
        return self.get_labels(self.alphabets, padding=False)


class Logger(object):
    _handle = None
    _root = None

    @staticmethod
    def init(output_dir, name, phase):
        format = '[%(asctime)s %(filename)s:%(lineno)d %(levelname)s {}] ' \
                '%(message)s'.format(name)
        logging.basicConfig(level=logging.INFO, format=format)

        try: os.makedirs(output_dir)
        except: pass
        config_path = os.path.join(output_dir, f'{phase}.txt')
        Logger._handle = logging.FileHandler(config_path)
        Logger._root = logging.getLogger()

    @staticmethod
    def enable_file():
        if Logger._handle is None or Logger._root is None:
            raise Exception('Invoke Logger.init() first!')
        Logger._root.addHandler(Logger._handle)

    @staticmethod
    def disable_file():
        if Logger._handle is None or Logger._root is None:
            raise Exception('Invoke Logger.init() first!')
        Logger._root.removeHandler(Logger._handle)


class Config(object):

    def __init__(self, config_path, host=True):
        def __dict2attr(d, prefix=''):
            for k, v in d.items():
                if isinstance(v, dict):
                    __dict2attr(v, f'{prefix}{k}_')
                else:
                    if k == 'phase':
                        assert v in ['train', 'test']
                    if k == 'stage':
                        assert v in ['pretrain-vision', 'pretrain-language',
                                     'train-semi-super', 'train-super']
                    self.__setattr__(f'{prefix}{k}', v)

        assert os.path.exists(config_path), '%s does not exists!' % config_path
        with open(config_path) as file:
            config_dict = yaml.load(file, Loader=yaml.FullLoader)
            #print("tiaoshi")
            #print(config_dict) #调试
            #print("config_dict after")#调试
        with open('/home/data4/zyh/ABINet/ABINet_Mindspore/configs/template.yaml') as file:
            default_config_dict = yaml.load(file, Loader=yaml.FullLoader)
            #print(default_config_dict) #调试
            #print("default_config_dict  after")
        __dict2attr(default_config_dict)
        __dict2attr(config_dict)
        self.global_workdir = os.path.join(self.global_workdir, self.global_name)

    def __getattr__(self, item):
        attr = self.__dict__.get(item)
        if attr is None:
            attr = dict()
            prefix = f'{item}_'
            for k, v in self.__dict__.items():
                if k.startswith(prefix):
                    n = k.replace(prefix, '')
                    attr[n] = v
            return attr if len(attr) > 0 else None
        else:
            return attr

    def __repr__(self):
        str = 'ModelConfig(\n'
        for i, (k, v) in enumerate(sorted(vars(self).items())):
            str += f'\t({i}): {k} = {v}\n'
        str += ')'
        return str

def blend_mask(image, mask, alpha=0.5, cmap='jet', color='b', color_alpha=1.0):
    # normalize mask
    mask = (mask-mask.min()) / (mask.max() - mask.min() + np.finfo(float).eps)
    if mask.shape != image.shape:
        mask = cv2.resize(mask,(image.shape[1], image.shape[0]))
    # get color map
    color_map = plt.get_cmap(cmap)
    mask = color_map(mask)[:,:,:3]
    # convert float to uint8
    mask = (mask * 255).astype(dtype=np.uint8)

    # set the basic color
    basic_color = np.array(colors.to_rgb(color)) * 255 
    basic_color = np.tile(basic_color, [image.shape[0], image.shape[1], 1]) 
    basic_color = basic_color.astype(dtype=np.uint8)
    # blend with basic color
    blended_img = cv2.addWeighted(image, color_alpha, basic_color, 1-color_alpha, 0)
    # blend with mask
    blended_img = cv2.addWeighted(blended_img, alpha, mask, 1-alpha, 0)

    return blended_img

def onehot(label, depth, device=None):
    """ 
    Args:
        label: shape (n1, n2, ..., )
        depth: a scalar

    Returns:
        onehot: (n1, n2, ..., depth)
    """
    # if not isinstance(label, ms.Tensor):
    #     label = ms.tensor(label)
    #onehot = torch.zeros(label.size() + torch.Size([depth]), device=device)

    #onehot_size = ms.Tensor.shape(depth) #好像是property类型的，没法调用


    # label_shape= label.shape
    # label_shape = np.int0(label_shape)
    # label_shape = label_shape[0]
    label_shape = 26


    #label_shape = np.int32(label_shape)
    # onehot_size = ms.ops.shape([depth])
    # onehot_size = label.shape + onehot_size
    #ms_onehot_zeros = ms.ops.Zeros()
    #ms_onehot_zeros = np.zeros()
    
    onehot_output = np.zeros((label_shape,depth)) #[26,37]的全0Tensor
    
    #label_expanddims = ms.ops.ExpandDims()
    #label_numpy_expanddims = np.expand_dims()
    label_expand = np.expand_dims(label,-1)

    #onehot = onehot.scatter_(-1, label.unsqueeze(-1), 1) #pytorch利用scatter_去创建onehot编码。
                                                #照scatter_各参数的意思，dim为-1(1),label.unsqueeze(-1)为index，1为src
                                                #原来onehot编码的shape为[26,37]，label.unsqueeze(-1)之后shape为[26,1]
    #onehot = ms.ops.scatter_nd(-1, label_expand, 1) Mindspore的scatter不能直接用，比较复杂，换了一个onehot算子，可能存在问题

    # updates = 
    # x = Tensor(np.array([[-0.1, 0.3, 3.6], [0.4, 0.5, -3.2]]).astype('float32'))
    # indices = Tensor(np.array([[0, 0], [0, 0]]).astype('int32'))
    # updates = Tensor(np.array([1.0, 2.2]).astype('float32'))



    # onehot_output = ms.Tensor(onehot_output)
    # label_expand = ms.Tensor(label_expand)
    # updates = ms.Tensor(np.ones((26,37)))
    # onehot_output = onehot_output.scatter_add(label_expand, updates)
    # onehot_output = np.zeros((4,5),np.float32) 


    label_expand = np.expand_dims(label_expand,-1)
    # print("label_expand.shape")
    # print(label_expand.shape)
    # print("label_expand")
    # print(label_expand)
    #onehot_output = ms.Tensor(onehot_output,ms.float32)
    
    #label_expand_onehot = ms.Tensor(np.random.rand(26,37))
    label_expand_onehot = np.zeros((26,37))
    a = 0
    for i in label_expand:
        i = int(i)
        
        #label_expand_onehot[i][0] = 1
        label_expand_onehot[a][i] = 1
        a = a + 1
    #updates = ms.Tensor(np.ones((26,37)),ms.float32)
    label_expand_onehot = label_expand_onehot

    onehot_output = label_expand_onehot + onehot_output





    # depth, on_value, off_value = 37, ms.Tensor(1.0, ms.float32), ms.Tensor(0.0, ms.float32)
    # onehot_output = ms.ops.one_hot(label ,depth, on_value, off_value, axis=-1)
    
    return onehot_output
